from .Dataset import Dataset
from dcr_reid.utils.utils import measure_time
from dcr_reid.utils.distance import global_dist
from dcr_reid.utils.distance import spatial_dist
from dcr_reid.utils.distance import normalize
from collections import defaultdict
from dcr_reid.utils.metric import cmc, mean_ap
from dcr_reid.utils.re_ranking import re_ranking
from dcr_reid.utils.dataset_utils import parse_im_name
import numpy as np
from PIL import Image
import time
import sys
import os.path as osp
class TestSet(Dataset):
    def __init__(self,im_dir=None,im_names=None,marks=None,extract_feat_func=None,separate_camera_set=None,single_gallery_shot=None,first_match_break=None,**kwargs):
        super(TestSet,self).__init__(dataset_size=len(im_names),**kwargs)
        self.im_dir = im_dir
        self.im_names = im_names
        self.marks = marks
        self.extract_feat_func = extract_feat_func
        self.separate_camera_set = separate_camera_set
        self.single_gallery_shot = single_gallery_shot
        self.first_match_break = first_match_break

    def set_feat_func(self, extract_feat_func):
        self.extract_feat_func = extract_feat_func

    def get_sample(self, ptr):
        im_name = self.im_names[ptr]
        im_path = osp.join(self.im_dir, im_name)
        im = np.asarray(Image.open(im_path))
        im, _ = self.pre_process_im(im)
        id = parse_im_name(self.im_names[ptr], 'id')
        cam = parse_im_name(self.im_names[ptr], 'cam')
        # denoting whether the im is from query, gallery, or multi query set.
        mark = self.marks[ptr]
        return im, id, cam, im_name, mark
    def next_batch(self):
        if self.epoch_done and self.shuffle:
            self.prng.shuffle(self.im_names)
        samples, self.epoch_done = self.prefetcher.next_batch()
        im_list, ids, cams, im_names, marks = zip(*samples)
        # Transform the list into a numpy array with shape[N,...]
        ims = np.stack(im_list, axis=0)
        ids = np.stack(ids)
        cams = np.array(im_names)
        im_names = np.array(im_names)
        marks = np.array(marks)
        return ims, ids, cams, im_names, marks, self.epoch_done

    def extract_feat(self, normalize_feat, verbose):
        '''
        Extract the features of the whole image set.
        :param normalize_feat: True or False, whether to normalize global and spatial feature to unit length
        :return:
        global_feats:numpy array with shape[N, C]
        spatial_feats: numpy array with shape[N,H,C]
        ids: numpy array with shape[N]
        cams:numpy array with shape[N]
        im_names: numpy array with shape[N]
        marks: numpy array with shape[N]
        '''
        global_feats, spatial_feats, ids, cams, im_names, marks = [],[],[],[],[],[]
        done = False
        step = 0
        printed = False
        st = time.time()
        last_time = time.time()
        while not done:
            ims_, ids_, cams_, im_names_, marks_, done_ = self.next_batch()
            global_feat, spatial_feat = self.extract_feat_func(ims_)
            global_feats.append(global_feat)
            spatial_feats.append(spatial_feat)
            ids.append(ids_)
            cams.append(cams_)
            im_names.append(im_names_)
            marks.append(marks_)
            #log, print the progress of extracting feature.
            if verbose:
                total_batches = (self.prefetcher.dataset_size//self.prefetcher.batch_size + 1)
                step+=1
                if step % 20 ==0:
                    if not printed:
                        printed = True
                    else:
                        # clean the current line
                        sys.stdout.write("\033[F\033[K")
                    print('{}/{} batches done, +{:.2f}s, total {:.2f}s'.format(step, total_batches,time.time()-last_time, time.time()-st))
                    last_time = time.time()
        global_feats = np.vstack(global_feats)
        spatial_feats = np.concatenate(spatial_feats)
        ids = np.hstack(ids)
        cams = np.hstack(cams)
        im_names = np.hstack(im_names)
        marks = np.hstack(marks)
        if normalize_feat:
            global_feats = normalize(global_feats, axis=1)
            spatial_feats = normalize(spatial_feats, axis=-1)
        return global_feats, spatial_feats, ids, cams, im_names, marks


    def eval(self, normalize_feat=True, to_re_rank=True, pool_type='average', verbose=True):
        '''
        Evaluate using metric CMC and mAP.
        Argus:
        :param normalize_feat: whether to normalize features before computing distance
        :param to_re_rank: whether to also report re-ranking scores.
        :param pool_type: 'average' or 'max', only for multi-query case.
        :param verbose: whether to print the intermediate information.
        '''
        with measure_time('Extracting feature...', verbose=verbose):
            global_feats, spatial_feats, ids, cams, im_names, marks = self.extract_feat(normalize_feat, verbose)
            # query, gallery, multi-query indices
            q_inds = marks == 0
            g_inds = marks == 1
            mq_inds = marks == 2

        def compute_score(dist_mat):
            mAp, cmc_scores = self.eval_map_cmc(q_g_dist=dist_mat,q_ids=ids[q_inds],g_ids=ids[g_inds],separate_camera_set=self.separate_camera_set,
                                                single_gallery_shot=self.single_gallery_shot,first_match_break=self.first_match_break,topk=10)
            return mAp, cmc_scores

        ###################
        # Single Query #
        ###################
        with measure_time('Single Query, Computing Distance...', verbose=verbose):
            # query-gallery distance using global distance.
            global_q_g_dist = global_dist(global_feats[q_inds], global_feats[g_inds], type='euclidean')
            spatial_q_g_dist = spatial_dist(spatial_feats[q_inds], spatial_feats[g_inds])
            for lam in range(11):
                weight = lam * 0.1
                with measure_time('Single Query, Computing scores...', verbose=verbose):
                    q_g_dist = (1-weight)*global_q_g_dist + weight*spatial_q_g_dist
                    print('{:<30}'.format('Single Query:'), end='')
                    mAP, cmc_scores = compute_score(q_g_dist)
        if to_re_rank:
            for lam in range(11):
                weight = lam * 0.1
                with measure_time('Single Query, Re-ranking...'):
                    #query-query distance.
                    global_q_q_dist = global_dist(global_feats[q_inds], global_feats[q_inds], type='euclidean')
                    spatial_q_q_dist = spatial_dist(spatial_feats[q_inds], spatial_feats[q_inds])
                    global_spatial_q_q_dist = (1-weight)*global_q_q_dist + weight*spatial_q_q_dist
                    # gallery-gallery distance.
                    global_g_g_dist = global_dist(global_feats[g_inds], global_feats[g_inds], type='euclidean')
                    spatial_g_g_dist = spatial_dist(spatial_feats[g_inds], spatial_feats[g_inds])
                    global_spatial_g_g_dist = (1-weight)*global_g_g_dist + weight*spatial_g_g_dist
                    #query-gallery distance.
                    global_spatial_q_g_dist = (1-weight)*global_q_g_dist + weight*spatial_q_g_dist
                    #re-ranked global query-gallery distance
                    re_r_global_q_g_dist = re_ranking(global_spatial_q_g_dist, global_spatial_q_q_dist, global_spatial_g_g_dist)
                with measure_time('Single Query, Computing scores for re-ranked Global Distance...'):
                    mAP, cmc_scores = compute_score(re_r_global_q_g_dist)

        ###################
        # Multi Query #
        ###################
        if any(mq_inds):
            mq_ids = ids[mq_inds]
            mq_cams = cams[mq_inds]
            mq_global_feat = global_feats[mq_inds]
            mq_spatial_feat = spatial_feats[mq_inds]
            unique_mq_ids_cams = defaultdict(list)
            for ind, (id, cam) in enumerate(zip(mq_ids, mq_cams)):
                unique_mq_ids_cams[(id, cam)].append(ind)
            keys = unique_mq_ids_cams.keys()
            assert pool_type in ['average', 'max']
            pool = np.mean if pool_type=='average' else np.max
            mq_global_feat = np.stack([pool(mq_global_feat[unique_mq_ids_cams[k]], axis=0) for k in keys])
            mq_spatial_feat = np.stack([pool(mq_spatial_feat[unique_mq_ids_cams[k]], axis=0) for k in keys])
            with measure_time('Multi Query, Computing distance...', verbose=verbose):
                # multi_query-gallery distance
                m_global_q_g_dist = global_dist(mq_global_feat, global_feats[g_inds], type='euclidean')
                m_spatial_q_g_dist = spatial_dist(mq_spatial_feat, spatial_feats[g_inds])
            for lam in range(11):
                weight = lam * 0.1
                with measure_time('Multi Query, Computing scores...', verbose=verbose):
                    m_q_g_dist = (1-weight)*m_global_q_g_dist + weight*m_spatial_q_g_dist
                    print('{:<30}'.format('Multi Query:'), end='')
                    mq_mAP, mq_cmc_scores = compute_score(m_q_g_dist, query_ids=np.array(zip(*keys)[0]),
                                                          gallery_ids=ids[g_inds], query_cams=np.array(zip(*keys)[1]),
                                                          gallery_cams=cams[g_inds])



    @staticmethod
    def eval_map_cmc(q_g_dist,q_ids=None,g_ids=None,q_cams=None,g_cams=None,separate_camera_set=None,
                     single_gallery_shot=None,first_match_break=None,topk=None):
        '''
        Compute CMC and mAP.
        :param q_g_dist: numpy array with shape[num_query, num_gallery], the pairwise distanced between query and gallery samples.
        :return:
        mAP: numpy array with shape[num_query], the AP averaged across query samples.
        cmc_scores: numpy array with shape[topk], the cmc curve averaged across query samples.
        '''
        # compute mean AP
        mAP = mean_ap(distmat=q_g_dist, query_ids=q_ids, gallery_ids=g_ids, query_cams=q_cams, gallery_cams=g_cams)
        # compute CMC scores
        cmc_scores  =cmc(distmat=q_g_dist, query_ids=q_ids, gallery_ids=g_ids, query_cams=q_cams, gallery_cams=g_cams,
                         separate_camera_set=separate_camera_set, single_gallery_shot=single_gallery_shot,
                         first_match_break=first_match_break, topk=topk)
        print('[mAP: {5.2%}], [rank1: {5.2%}], [rank3: {5.2%}], [rank5: {5.2%}], [rank10: {5.2%}]'.format(mAP, *cmc_scores[0,2,4,9]))
        return mAP, cmc_scores
